﻿#include "TaylorGreenVortex.h"
#include "Algorithm.h"
#include "DataObject.h"
#include "Mesh.h"
#include <vtkm/VectorAnalysis.h>

#include <vtkm/worklet/WorkletMapField.h>
#include <vtkm/worklet/WorkletMapTopology.h>

#include <vtkm/Math.h>
#include <vtkm/cont/ArrayHandleGroupVec.h>
#include <vtkm/exec/ParametricCoordinates.h>

struct UpdateVelocityAndPosition1 : vtkm::worklet::WorkletMapField
{
  using ControlSignature = void(FieldIn bnd,
                                FieldIn point_mass,
                                FieldInOut force,
                                FieldInOut velocity,
                                FieldInOut position);
  using ExecutionSignature = void(_1, _2, _3, _4, _5);

  UpdateVelocityAndPosition1(Real dt)
    : _dt(dt)
  {
  }
  template<typename CoordType>
  VTKM_EXEC void operator()(const Id& bid,
                            const Real& point_mass,
                            Vec3f& force,
                            Vec3f& vel,
                            CoordType& point) const
  {
   
    //if (bid == 0 || bid == 2)
    //{
    //  force[1] = 0;
    //}
    //else if (bid == 1 || bid == 3)
    //{
    //  force[0] = 0;
    //}
    //else if (bid == 4)
    //{
    //  force[0] = 0;
    //  force[1] = 0;
    //}
    vel += force * _dt / point_mass;
    point += vel * _dt;
  }

  Real _dt;
};

struct UpdateCellEnergyWithSource : WorkletVisitCellsWithPoints
{
  // clang-format off
  using ControlSignature =  void(CellSetIn, FieldInCell mass, FieldInPoint velocity, FieldInCell force, FieldInOutCell energy,  FieldInPoint pts);
  // clang-format on
  using ExecutionSignature = void(PointCount, _2, _3, _4, _5, _6);

  UpdateCellEnergyWithSource(Real dt)
    : _dt(dt)
  {
  }
  template<typename ForceType, typename VelocityType, typename PType>
  VTKM_EXEC void operator()(const IdComponent& num_pts,
                            const Real& mass,
                            const VelocityType& velocity,
                            const ForceType& force,
                            Real& energy,
                            const PType& pts) const
  {
    Real rhs = 0;
    for (IdComponent i = 0; i < num_pts; i++)
    {
      rhs += vtkm::Dot(force[i], velocity[i]);
    }
    auto p = (pts[0] + pts[1] + pts[2]) / 3;

    auto x = p[0];
    auto y = p[1];
    constexpr auto pi = vtkm::Pi();
    auto source = 3. / 8 * pi *
      (vtkm::Cos(3 * pi * x) * vtkm::Cos(pi * y) - vtkm::Cos(pi * x) * vtkm::Cos(3 * pi * y));
    //source = 0;
    energy += _dt * (-rhs / mass + source);
  }

  Real _dt;
};

struct InitialConditionOnCell1 : vtkm::worklet::WorkletVisitCellsWithPoints
{
  using ValueType = Vec<FloatDefault, 5>;

  using ControlSignature = void(CellSetIn, FieldInPoint, FieldOut density, FieldOut pressure);
  using ExecutionSignature = void(CellShape, _2, _3, _4);

  template<typename CellShapeTag, typename PType>
  VTKM_EXEC void operator()(const CellShapeTag& shape,
                            const PType& pts,
                            FloatDefault& density,
                            FloatDefault& pressure) const
  {
    auto p = (pts[0] + pts[1] + pts[2]) / 3;

    auto x = p[0];
    auto y = p[1];
    constexpr auto pi = vtkm::Pi();
    density = 1;
    pressure = 1. / 4 * (vtkm::Cos(2 * pi * x) + vtkm::Cos(2 * pi * y)) + 1;
  }
};

struct InitialConditionOnPoint1 : vtkm::worklet::WorkletMapField
{
  using ControlSignature = void(FieldIn point, FieldOut velocity);
  using ExecutionSignature = void(_1, _2);

  VTKM_EXEC
  void operator()(const Vec3f& p, Vec3f& velocity) const
  {
    constexpr auto pi = vtkm::Pi();
    auto x = p[0];
    auto y = p[1];
    Real u = vtkm::Sin(pi * x) * vtkm::Cos(pi * y);
    Real v = -vtkm::Cos(pi * x) * vtkm::Sin(pi * y);
    velocity = Vec3f{ u, v, 0 };
  }
};


struct ComputePointForceWithSlipBC : vtkm::worklet::WorkletReduceByKey
{
  struct ValuesInWholeArray : vtkm::cont::arg::ControlSignatureTagBase
  {
    using TypeCheckTag = vtkm::cont::arg::TypeCheckTagArrayIn;
    using TransportTag = vtkm::cont::arg::TransportTagKeyedValuesIn;
    using FetchTag = vtkm::exec::arg::FetchTagExecObject;
  };

  using ControlSignature = void(KeysIn,
                                ValuesIn c2n_force,
                                ReducedValuesOut point_force,
                                ValuesInWholeArray,
                                ReducedValuesIn bid,
                                ReducedValuesIn periodic_map);
  using ExecutionSignature = void(_1, _2, _3, _4, _5, _6, WorkIndex);

  template<typename ValueType, typename WholeArray>
  VTKM_EXEC void operator()(const Id& key,
                            const ValueType& c2n_force,
                            Vec3f& point_force,
                            WholeArray& arr,
                            const Id& bid,
                            const Id& map_id,
                            Id& tid) const

  {
    point_force = Vec3f{ 0, 0, 0 };
    auto& force = arr.Get(key);
    for (IdComponent i = 0; i < force.GetNumberOfComponents(); i++)
    {
      point_force += force[i];
    }

    // 附加边界条件
    if (bid == 0 || bid == 2)
    {
      point_force[1] = 0;
    }
    else if (bid == 1 || bid == 3)
    {
      point_force[0] = 0;
    }
    else if (bid == 4)
    {
      point_force[0] = 0;
      point_force[1] = 0;
    }
  }
};

TaylorGreenVortex::TaylorGreenVortex(const Configuration& cfg)
  : SGHSystem(cfg)
{
}

RegisterObject(TaylorGreenVortex);
void TaylorGreenVortex::InitialCondition()
{
  auto& cellset = _data_set.GetCellSet();
  auto& coords = _data_set.GetCoordinateSystem();
  auto& density = GetVariableArray<Real>("density");
  auto& pressure = GetVariableArray<Real>("pressure");
  auto& vel = GetVariableArray<Vec3f>("velocity");

  Invoker{}(InitialConditionOnCell1{}, cellset, coords, density, pressure);
  Invoker{}(InitialConditionOnPoint1{}, coords, vel);
}

void TaylorGreenVortex::PreSolve()
{
  UpdateCellToPointForce();
  UpdatePointForce();
}

void TaylorGreenVortex::UpdatePointForce()
{
  auto& force = GetVariableArray<Vec3f>("force");
  auto& bid = _mesh.GetBoundaryId();
  vtkm::worklet::Keys<Id> keys(_c2n);
  invoker(ComputePointForceWithSlipBC{}, keys, _c2n_force, force, _c2n_force, bid, _periodic_map);
}

void TaylorGreenVortex::Solve()
{
  ComputeTimeStepper(); // 计算 dt

  auto& point_mass = GetVariableArray<Real>("point_mass");
  auto& force = GetVariableArray<Vec3f>("force");
  auto& velocity = GetVariableArray<Vec3f>("velocity");
  auto& position = _mesh.GetCoordinate();
  auto& bnd_id = _mesh.GetBoundaryId();

  invoker(UpdateVelocityAndPosition1{ _dt }, bnd_id, point_mass, force, velocity, position);

  auto cell_mass = GetVariableArray<Real>("cell_mass");
  auto energy = GetVariableArray<Real>("energy");

  invoker(UpdateCellEnergyWithSource{ _dt },
          _mesh.GetCellSet(),
          cell_mass,
          velocity,
          vtkm::cont::make_ArrayHandleGroupVec<3>(_c2n_force),
          energy,
          _mesh.GetCoordinate());
}
